import ijson
import os
import sys
import numpy as np
from PyCmpltrtok.common import sep


def analyze_data(xdir, xdir_squad, name):
    sep(name)
    
    path = os.path.join(xdir, name)
    print(f'path=|{path}|')
    p, q, a = 0, 0, 0
    pn, qn, an = 0, 0, 0
    xas = []
    with open(path, 'rb') as f:
        xitems_gen = ijson.items(f, 'item')
        for xrecord in xitems_gen:
            context = xrecord.get('context_text', '')
            if context:
                p += 1
            else:
                pn += 1
            qas = xrecord.get('qas', [])
            for qa in qas:
                question = qa.get('query_text', '')
                if question:
                    q += 1
                else:
                    qn += 1
                answers = qa.get('answers', [])
                xas_i = 0
                for ans in answers:
                    if ans:
                        a += 1
                        xas_i += 1
                    else:
                        an += 1
                xas.append(xas_i)
    xas = np.array(xas, dtype='f2')
    print(f'p: {p}, q: {q}, a: {a}, a per q: {xas.mean()}')
    print(f'pn: {pn}, qn: {qn}, an: {an}')
    
    path = os.path.join(xdir_squad, name)
    print(f'path=|{path}|')
    p, q, a = 0, 0, 0
    pn, qn, an = 0, 0, 0
    xas = []
    with open(path, 'rb') as f:
        xitems_gen = ijson.items(f, 'data.item')
        for xrecord in xitems_gen:
            paragraphs = xrecord.get('paragraphs', [])
            for para in paragraphs:
                context = para.get('context', '')
                if context:
                    p += 1
                else:
                    pn += 1
                qas = para.get('qas', [])
                for qa in qas:
                    question = qa.get('question', '')
                    if question:
                        q += 1
                    else:
                        qn += 1
                    answers = qa.get('answers', [])
                    xas_i = 0
                    for ans in answers:
                        text = ans.get('text', '')
                        if text:
                            a += 1
                            xas_i += 1
                        else:
                            an += 1
                    xas.append(xas_i)
    xas = np.array(xas, dtype='f2')
    print(f'p: {p}, q: {q}, a: {a}, a per q: {xas.mean()}')
    print(f'pn: {pn}, qn: {qn}, an: {an}')


if '__main__' == __name__:
    
    def main():
        # xdir = '/home/peiyp2004/code_github/cmrc2018/data/'  # WSL
        xdir = '/home/yunpeng/code_github/cmrc2018/data/'  # WSL
        # xdir_squad = '/home/peiyp2004/code_github/cmrc2018/squad-style-data/'  # WSL
        xdir_squad = '/home/yunpeng/code_github/cmrc2018/squad-style-data/'  # New PC
        names = [
            'cmrc2018_train.json',
            'cmrc2018_dev.json',
            'cmrc2018_trial.json',
        ]
        
        for name in names:
            analyze_data(xdir, xdir_squad, name)
            
    
    main()
    sep('All over')
    